using System.Collections.Generic;
using Latios.Authoring;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Entities;
using Unity.Mathematics;
using UnityEngine;

namespace Latios.Unika.Authoring
{
    /// <summary>
    /// A base class for authoring an untyped script. Prefer to use UnikaScriptAuthoring<T> instead unless you need to do something very custom.
    /// However, you should still refer to this base type to understand the methods you need to override.
    /// </summary>
    [RequireComponent(typeof(UnikaScriptBufferAuthoring))]
    public abstract class UnikaScriptAuthoringBase : MonoBehaviour, IUnikaAuthoringScriptCounter
    {
        private protected static List<IUnikaAuthoringScriptCounter> m_scriptsCache = new List<IUnikaAuthoringScriptCounter>();

        /// <summary>
        /// Implement this method to determine whether this authoring component will produce a valid script.
        /// This method will be called frequently from various points in baking. It is crucial that this method
        /// returns false if the script baking will be canceled for any reason, as otherwise ScriptRefs obtained
        /// from any baked scripts on the target entity at bake time may be corrupted.
        /// </summary>
        /// <returns>True if a script will be baked given the current configuration, false otherwise</returns>
        public abstract bool IsValid();

        /// <summary>
        /// Implement this method to control the baking of the script.
        /// </summary>
        /// <param name="baker">A baker used to generate the script</param>
        /// <param name="toAssign">An object to assign the script data to, as well as configure the user byte and flags</param>
        /// <param name="smartPostProcessTarget">A baking-only entity from which you can receive a reference to the script in an ISmartPostProcessItem
        /// before the script is merged into the scripts buffer</param>
        public abstract void Bake(IBaker baker, ref AuthoredScriptAssignment toAssign, Entity smartPostProcessTarget);

        /// <summary>
        /// Gets the untyped ScriptRef for this script. Call this from other authoring scripts or bakers
        /// </summary>
        /// <param name="baker">The baker being used to bake whatever consumes the ScriptRef</param>
        /// <param name="transformUsageFlags">The transform flags that should be added to the script's entity, if any</param>
        /// <returns>An untyped ScriptRef for the script generated by this authoring component, or a Null ScriptRef if this authoring is not valid.</returns>
        public ScriptRef GetScriptRef(IBaker baker, TransformUsageFlags transformUsageFlags = TransformUsageFlags.None)
        {
            if (!enabled)
                return default;

            m_scriptsCache.Clear();
            baker.GetComponents(this, m_scriptsCache);
            for (int i = 0, validBefore = 0; i < m_scriptsCache.Count; i++)
            {
                var s = m_scriptsCache[i];
                if (ReferenceEquals(s, this))
                {
                    if (IsValid())
                    {
                        return new ScriptRef
                        {
                            m_entity            = baker.GetEntity(this, transformUsageFlags),
                            m_instanceId        = validBefore + 1,
                            m_cachedHeaderIndex = validBefore
                        };
                    }
                    else
                        return default;
                }
                else
                    validBefore += s.CountScripts();
            }
            return default;
        }

        /// <summary>
        /// Assigns the Script and extracts a typed ScriptRef which can be used to assign references to other components being baked
        /// </summary>
        /// <typeparam name="T">The type of script to be generated</typeparam>
        /// <param name="baker">The baker being used to bake this script</param>
        /// <param name="script">The script field data the script should be initialized with</param>
        /// <param name="assignment">The assignment object to assign the script data to</param>
        /// <param name="transformUsageFlags">The transform flags that should be added to the script's entity, if any</param>
        /// <returns>A typed ScriptRef for the assigned script</returns>
        protected ScriptRef<T> AssignAndGetTypedRef<T>(IBaker baker,
                                                       ref T script,
                                                       ref AuthoredScriptAssignment assignment,
                                                       TransformUsageFlags transformUsageFlags = TransformUsageFlags.None)
            where T : unmanaged, IUnikaScript, IUnikaScriptGen
        {
            assignment.Assign(ref script);
            var scriptRef = GetScriptRef(baker, transformUsageFlags);
            return new ScriptRef<T>
            {
                m_entity            = scriptRef.m_entity,
                m_instanceId        = scriptRef.m_instanceId,
                m_cachedHeaderIndex = scriptRef.m_cachedHeaderIndex,
            };
        }

        int IUnikaAuthoringScriptCounter.CountScripts()
        {
            return (enabled && IsValid()) ? 1 : 0;
        }
    }

    /// <summary>
    /// A structure which represents a script being baked. Assign the script structure, userByte, userFlagA, and userFlagB to this structure.
    /// </summary>
    public struct AuthoredScriptAssignment
    {
        internal NativeArray<byte> scriptPayload;
        internal int               scriptType;
        /// <summary>
        /// The TransformUsageFlags that should be added to the entity containing all the scripts.
        /// </summary>
        public TransformUsageFlags transformUsageFlags;
        /// <summary>
        /// The userByte the script will initially have
        /// </summary>
        public byte userByte;
        /// <summary>
        /// The userFlagA value the script will initially have
        /// </summary>
        public bool userFlagA;
        /// <summary>
        /// The userFlagB value the script will initially have
        /// </summary>
        public bool userFlagB;

        /// <summary>
        /// Assigns the script field members using a local copy of the script structure
        /// </summary>
        /// <typeparam name="T">The type of script to be generated</typeparam>
        /// <param name="script">The script field data the script should be initialized with</param>
        public unsafe void Assign<T>(ref T script) where T : unmanaged, IUnikaScript, IUnikaScriptGen
        {
            var scriptSize = UnsafeUtility.SizeOf<T>();
            scriptType     = ScriptTypeInfoManager.GetScriptRuntimeIdAndMask<T>().runtimeId;
            scriptPayload  = new NativeArray<byte>(scriptSize, Allocator.Temp);
            UnsafeUtility.MemCpy(scriptPayload.GetUnsafePtr(), UnsafeUtility.AddressOf(ref script), scriptSize);
        }
    }

    public static class UnikaBakingUtilities
    {
        /// <summary>
        /// Adds a script to the specified object. Scripts added this way do not need to be added from an authoring component on the same GameObject they target.
        /// </summary>
        /// <typeparam name="T">The type of script to add</typeparam>
        /// <param name="targetBuffer">The GameObject for the entity this script should be added to. This does not need to be the same GameObject the authoring component belongs to.</param>
        /// <param name="script">The script data to be assigned</param>
        /// <param name="userByte">The userByte value the script should be initialized with</param>
        /// <param name="userFlagA">The userFlagA value the script should be initialized with</param>
        /// <param name="userFlagB">The userFlagB value the script should be initialized with</param>
        /// <returns>A baking-only entity from which you can receive a reference to the script in an ISmartPostProcessItem
        /// before the script is merged into the scripts buffer</returns>
        public static unsafe Entity AddScript<T>(this IBaker baker,
                                                 UnikaScriptBufferAuthoring targetBuffer,
                                                 ref T script,
                                                 byte userByte = 0,
                                                 bool userFlagA = false,
                                                 bool userFlagB = false) where T : unmanaged, IUnikaScript, IUnikaScriptGen
        {
            var entity = baker.CreateAdditionalEntity(TransformUsageFlags.None, true);
            var size   = UnsafeUtility.SizeOf<T>();
            baker.AddComponent(entity, new BakedScriptMetadata
            {
                scriptRef = new ScriptRef
                {
                    m_entity            = baker.GetEntity(targetBuffer, TransformUsageFlags.None),
                    m_cachedHeaderIndex = -1,
                    m_instanceId        = 0
                },
                scriptType = ScriptTypeInfoManager.GetScriptRuntimeIdAndMask<T>().runtimeId,
                userFlagA  = userFlagA,
                userByte   = userByte,
                userFlagB  = userFlagB
            });
            var bytes = baker.AddBuffer<BakedScriptByte>(entity);
            bytes.ResizeUninitialized(size);
            UnsafeUtility.MemCpy(bytes.GetUnsafePtr(), UnsafeUtility.AddressOf(ref script), size);
            return entity;
        }

        /// <summary>
        /// Retrieves the baked script by reference so that blob assets can be assigned from smart blobbers
        /// </summary>
        /// <typeparam name="T">The type of script</typeparam>
        /// <param name="bakingOnlyEntity">The baking-only entity that contains the script</param>
        /// <returns></returns>
        public static unsafe ref T GetBakedScriptInPostProcess<T>(this EntityManager entityManager, Entity bakingOnlyEntity) where T : unmanaged, IUnikaScript, IUnikaScriptGen
        {
            var bytes = entityManager.GetBuffer<BakedScriptByte>(bakingOnlyEntity);
            return ref UnsafeUtility.AsRef<T>(bytes.GetUnsafePtr());
        }

        /// <summary>
        /// Sets up an additional entity to bake scripts and have them be properly serialized in the subscene.
        /// To add scripts, use the ScriptStructuralChange API
        /// </summary>
        /// <param name="additionalEntity">The additional entity to add the components and buffers to</param>
        public static DynamicBuffer<UnikaScripts> CreateScriptBuffersForAdditionalEntity(this IBaker baker, Entity additionalEntity)
        {
            FixedList128Bytes<ComponentType> types = default;
            types.Add(ComponentType.ReadWrite<UnikaScripts>());
            types.Add(ComponentType.ReadWrite<UnikaSerializedEntityReference>());
            types.Add(ComponentType.ReadWrite<UnikaSerializedBlobReference>());
            types.Add(ComponentType.ReadWrite<UnikaSerializedAssetReference>());
            types.Add(ComponentType.ReadWrite<UnikaSerializedObjectReference>());
            types.Add(ComponentType.ReadWrite<UnikaSerializedTypeIds>());
            baker.AddComponent(additionalEntity, new ComponentTypeSet(in types));
            return baker.SetBuffer<UnikaScripts>(additionalEntity);
        }
    }

    [BakeDerivedTypes]
    public class UnikaScriptAuthoringBaseBaker : Baker<UnikaScriptAuthoringBase>
    {
        public override void Bake(UnikaScriptAuthoringBase authoring)
        {
            if (!authoring.IsValid())
                return;

            var                      entity     = CreateAdditionalEntity(TransformUsageFlags.None, true);
            AuthoredScriptAssignment assignment = default;
            authoring.Bake(this, ref assignment, entity);
            if (assignment.scriptPayload.Length == 0)
                throw new System.InvalidOperationException(
                    $"The AuthoredScriptAssignment was left unassigned for object {authoring.name}. Please fix this issue before attempting to enter play mode, or else Unity may crash.");

            GetEntity(assignment.transformUsageFlags);
            var bytes = AddBuffer<BakedScriptByte>(entity).Reinterpret<byte>();
            bytes.AddRange(assignment.scriptPayload);
            AddComponent(entity, new BakedScriptMetadata
            {
                scriptRef  = authoring.GetScriptRef(this),
                userFlagA  = assignment.userFlagA,
                scriptType = assignment.scriptType,
                userByte   = assignment.userByte,
                userFlagB  = assignment.userFlagB,
            });
        }
    }
}

